# n-snakes: Multidimentional, procedural exploration game.
# http://code.google.com/p/n-snakes/
#
#
# Copyright (c) 2012 Eliot Eshelman
#
# Built upon Snakes demo written by Jessica McKellar for OpenHatch
# https://openhatch.org/wiki/Snakes_Pygame
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.
###############################################################################

import pygame
import utils

snake_directions = utils.enum((
    "UP",
    "DOWN",
    "LEFT",
    "RIGHT"
))


class Snake_Segment(pygame.sprite.Sprite):
    """
    One segment of a Snake.

    Members:
        image  - visible pygame Surface
        rect   - the dimensions of the segment
        color  - the color of the segment
    """

    def __init__(self, color, pos, size):
        # Initialize parent Sprite class
        pygame.sprite.Sprite.__init__(self)

        self.image = pygame.Surface(size)
        self.color = color
        self.image.fill(self.color)

        self.rect = self.image.get_rect()
        self.rect.topleft = pos


class Eaten_Food(pygame.sprite.Sprite):
    """
    A piece of food that has been eaten by the snake.

    Members:
        image     - visible pygame Surface
        rect      - the dimensions of the food
        color     - the color of the food
        body_pos  - integer position within the snake's body
        digested  - timer for keeping track of digestion
    """

    def __init__(self, color, pos, size):
        # Initialize parent Sprite class
        pygame.sprite.Sprite.__init__(self)

        self.image = pygame.Surface(size)
        self.color = color
        self.image.fill(self.color)

        self.rect = self.image.get_rect()
        self.rect.topleft = pos

        self.body_pos = 0
        self.digested = pygame.time.get_ticks()


class Snake(object):
    """
    A snake. It has a head and body and can move.
    If the snake shrinks smaller than the head plus one segment, it's dead.

    Members:
        _head_group    - Sprite Group with the snake's head
        _segment_group - Sprite Group with the snake's segments
        _food_group    - Sprite Group with food the snake has eaten
        _color         - the color of the snake
        _direction     - which direction the snake is headed
    """

    SNAKE_UNIT = 20
    SEGMENT_SIZE = (SNAKE_UNIT, SNAKE_UNIT)

    EATEN_FOOD_UNIT = 10
    EATEN_FOOD_SIZE = (EATEN_FOOD_UNIT, EATEN_FOOD_UNIT)
    EATEN_FOOD_OFFSET = (SNAKE_UNIT - EATEN_FOOD_UNIT)/2
    FOOD_DIGEST_TIME = 1000  # 1 second

    def __init__(self, color, length=2, pos=[0, 0]):
        """
        Create snake head and body with (length-1) segments.
        """
        # Create the Sprite Groups
        self._head_group = pygame.sprite.RenderUpdates()
        self._segment_group = pygame.sprite.OrderedUpdates()
        self._food_group = pygame.sprite.RenderUpdates()

        # The body color is slightly lighter or darker than the head color.
        if (color.r + color.g + color.b) < 0xAA:
            self._color = utils.brighten_color(color, 0x30)
        else:
            self._color = utils.dim_color(color, 0x30)

        # Keep track of the direction of movement.
        self._direction = snake_directions.UP

        # Create the head.
        head = Snake_Segment(color, pos, self.SEGMENT_SIZE)
        self._head_group.add(head)

        # Create the body segments.
        for i in range(length - 1):
            self._add_segment()


    def _add_segment(self):
        """
        Add a segment to the body of the snake.
        """
        if self._segment_group:
            tail = self._segment_group.sprites()[-1]
        else:
            tail = self._head_group.sprites()[0]
        x, y = tail.rect.topleft

        # Add to the tail depending on the direction of movement.
        if self._direction == snake_directions.UP:
            y += self.SNAKE_UNIT
        elif self._direction == snake_directions.DOWN:
            y -= self.SNAKE_UNIT
        elif self._direction == snake_directions.LEFT:
            x += self.SNAKE_UNIT
        elif self._direction == snake_directions.RIGHT:
            x -= self.SNAKE_UNIT
        
        segment = Snake_Segment(self._color, (x, y), self.SEGMENT_SIZE)
        self._segment_group.add(segment)


    def _remove_segment(self):
        """
        Removes one segment from the tail of the snake.
        """
        # TODO
        pass


    def _add_eaten_food(self, food_color):
        """
        Add a piece of food the snake has eaten.
        """
        # Add the food just behind the head.
        segment = self._segment_group.sprites()[0]
        x, y = segment.rect.topleft
        x += self.EATEN_FOOD_OFFSET
        y += self.EATEN_FOOD_OFFSET

        food = Eaten_Food(food_color, (x, y), self.EATEN_FOOD_SIZE)
        self._food_group.add(food)


    def _update_food(self):
        """
        Update the eaten food in the snake.
        """
        current_time = pygame.time.get_ticks()

        for food in self._food_group.sprites():
            last_time = food.digested
            digested = False
            if (current_time - last_time) > self.FOOD_DIGEST_TIME:
                digested = self._digest_food(food, current_time)
            if not digested:
                pos = self._segment_group.sprites()[food.body_pos].rect.topleft
                food.rect.topleft = [
                        pos[0] + self.EATEN_FOOD_OFFSET,
                        pos[1] + self.EATEN_FOOD_OFFSET
                    ]


    def _digest_food(self, food, time):
        """
        When a snake consumes food, it passes through the body.

        Returns:
            False - if food is still being digested.
            True - if digestion is complete.
        """
        tail_length = len(self._segment_group)
        if food.body_pos >= (tail_length - 1):
            self._food_group.remove(food)
            
            # Add another segment to the tail.
            self._add_segment()
            return True
        else:
            food.body_pos += 1
            food.digested = time
            return False


    def head_hit_body(self):
        """
        Return True if the head hit the body.
        """
        head = self._head_group.sprites()[0]
        if pygame.sprite.spritecollideany(head, self._segment_group):
            return True
        else:
            return False


    def die(self):
        """
        Called when this snake has died.
        """
        for group in (self._head_group, self._segment_group, self._food_group):
            for sprite in group.sprites():
                sprite.kill()


    def get_head(self):
        """
        Return the head of the snake.
        """
        return self._head_group.sprites()[0]


    def get_all_segments(self):
        """
        Return all segments of the snake, including the head.
        """
        return self._head_group.sprites() + self._segment_group.sprites()


    def update(self):
        """
        Update the snake.
        """
        self._update_food()


    def draw(self, surface):
        """
        Draw the snake sprites.

        Return a list of the modified rectangles - for use by clear()
        """
        dirty_rects = self._head_group.draw(surface)
        dirty_rects += self._segment_group.draw(surface)
        dirty_rects += self._food_group.draw(surface)
        return dirty_rects


    def clear(self, surface, background):
        """
        Clear the snake sprites from the previous loop.
        """
        self._head_group.clear(surface, background)
        self._segment_group.clear(surface, background)
        self._food_group.clear(surface, background)


    def move(self, key, game_surface):
        """
        Move the snake.
        """
        game_rect = game_surface.get_rect()
        head = self._head_group.sprites()[0]
        old_direction = self._direction

        # Move the position of the head based on the provided keypress.
        # Refuse if the move would put the head backwards into the body.
        head_x, head_y = head.rect.topleft
        if (key == pygame.K_UP):
            if self._direction == snake_directions.DOWN:
                return
            else:
                head_y = head_y - self.SNAKE_UNIT
                self._direction = snake_directions.UP
        elif (key == pygame.K_DOWN):
            if self._direction == snake_directions.UP:
                return
            else:
                head_y = head_y + self.SNAKE_UNIT
                self._direction = snake_directions.DOWN
        elif (key == pygame.K_RIGHT):
            if self._direction == snake_directions.LEFT:
                return
            else:
                head_x = head_x + self.SNAKE_UNIT
                self._direction = snake_directions.RIGHT
        elif (key == pygame.K_LEFT):
            if self._direction == snake_directions.RIGHT:
                return
            else:
                head_x = head_x - self.SNAKE_UNIT
                self._direction = snake_directions.LEFT
        else:
            return

        # Don't actually move if that would put the head outside the
        # bounds of the screen.
        max_x = game_rect.right - self.SNAKE_UNIT
        max_y = game_rect.bottom - self.SNAKE_UNIT
        if head_x < 0 or head_x > max_x or head_y < 0 or head_y > max_y:
            self._direction = old_direction
            return

        # Update the head and body segments
        old_segment_pos = head.rect.topleft
        head.rect.topleft = (head_x, head_y)
        for segment in self._segment_group.sprites():
            saved_pos = segment.rect.topleft
            segment.rect.topleft = old_segment_pos
            old_segment_pos = saved_pos

        return True


    def try_to_eat(self, snake_food):
        """
        Determine if the snake's head has hit some food.
        """
        head = self._head_group.sprites()[0]

        # Remove the food sprites that have been eaten
        for f in pygame.sprite.spritecollide(head, snake_food, True):
            self._add_eaten_food(f.color)
